import pandas as pd
from sklearn.model_selection import GroupKFold
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
import shap
import matplotlib.pyplot as plt
import os
import joblib

# 读取数据
data_path = 'D:\\学习&科研\\华为手表项目\\华为数据\\试验记录表\\all_stages_df_statistics.csv'
df = pd.read_csv(data_path)

# 筛选出 state 为 'running' 的行
df = df[df['state'] == 'running']

# 将 polar_hr 和 polar_rr 列转换为适合模型的格式，例如取平均值
df['polar_hr'] = df['polar_hr'].apply(lambda x: eval(x)[0] if isinstance(eval(x), list) and len(eval(x)) > 0 else 0)
df['polar_rr'] = df['polar_rr'].apply(lambda x: eval(x)[0] if isinstance(eval(x), list) and len(eval(x)) > 0 else 0)

# 定义输入特征和目标变量
X = df[['la', 'speed', 'polar_hr_mean', 'polar_hr_min', 'polar_hr_max', 'polar_hr_median',
         'polar_hr_q1', 'polar_hr_q3', 'polar_rr_mean', 'polar_rr_median',
         'polar_rr_q1', 'polar_rr_q3', 'sex', 'age', 'hight', 'weight']]
y = df['physiology_RPE']
groups = df['number']  # 分组列

# 使用 GroupKFold 进行分组划分
group_kfold = GroupKFold(n_splits=10)

# 初始化结果列表
results = []
all_y_true = []
all_y_pred = []

# 创建保存SHAP图和比较图的文件夹
shap_plot_dir = 'shap_plots'
os.makedirs(shap_plot_dir, exist_ok=True)

# 基于分组进行划分和训练
for fold, (train_index, test_index) in enumerate(group_kfold.split(X, y, groups=groups)):
    X_train, X_test = X.iloc[train_index], X.iloc[test_index]
    y_train, y_test = y.iloc[train_index], y.iloc[test_index]

    # 创建随机森林回归模型
    model = RandomForestRegressor(n_estimators=100, random_state=42)

    # 训练模型
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)

    # 评估模型
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    mae = mean_absolute_error(y_test, y_pred)

    # 保存每个分组的结果
    results.append({
        'Mean Squared Error': mse,
        'R² Score': r2,
        'Mean Absolute Error': mae
    })

    # 保存所有的真实值和预测值
    all_y_true.extend(y_test)
    all_y_pred.extend(y_pred)

    # 计算 SHAP 值
    explainer = shap.TreeExplainer(model)  # 使用 TreeExplainer
    shap_values = explainer.shap_values(X_test)

    # 绘制 SHAP 值图
    plt.figure()
    shap.summary_plot(shap_values, X_test, plot_type="bar", show=False)
    plt.title(f'SHAP Summary Plot - Fold {fold + 1}')
    plt.savefig(os.path.join(shap_plot_dir, f'shap_summary_fold_{fold + 1}.png'))  # 保存图像
    plt.close()

# 将结果转换为 DataFrame 并输出
results_df = pd.DataFrame(results)
print(results_df)

# 计算整体模型性能
overall_mse = mean_squared_error(all_y_true, all_y_pred)
overall_r2 = r2_score(all_y_true, all_y_pred)
overall_mae = mean_absolute_error(all_y_true, all_y_pred)

# 输出整体模型性能
print(f'Overall Mean Squared Error: {overall_mse}')
print(f'Overall R² Score: {overall_r2}')
print(f'Overall Mean Absolute Error: {overall_mae}')

# 使用最佳模型在整个训练集上进行最终训练
final_model = RandomForestRegressor(n_estimators=100, random_state=42)
final_model.fit(X, y)

# 保存最终模型
joblib.dump(final_model, 'final_random_forest_model.joblib')  # 保存模型

# 加载模型并进行预测
loaded_model = joblib.load('final_random_forest_model.joblib')

# 进行预测（使用整个数据集或新的数据集）
predictions = loaded_model.predict(X)

# 创建一个 DataFrame 来保存真实值和预测值
comparison_df = pd.DataFrame({
    'True Values': y,
    'Predictions': predictions,
    'number': df['number']
})

# 按 number 保存每个 group's 预测结果，并绘制图形
for number in comparison_df['number'].unique():
    group_df = comparison_df[comparison_df['number'] == number]
    
    # 保存预测结果到 CSV 文件
    group_df.to_csv(f'predictions_comparison_random_forest_number_{number}.csv', index=False)
    
    # 绘制实际值和预测值的对比图
    plt.figure()
    plt.scatter(group_df['True Values'], group_df['Predictions'], color='blue', alpha=0.6)
    plt.plot([group_df['True Values'].min(), group_df['True Values'].max()],
             [group_df['True Values'].min(), group_df['True Values'].max()], color='red', linestyle='--')
    plt.title(f'True vs Predicted Values for Number {number}')
    plt.xlabel('True Values')
    plt.ylabel('Predicted Values')
    plt.xlim(group_df['True Values'].min() - 1, group_df['True Values'].max() + 1)
    plt.ylim(group_df['Predictions'].min() - 1, group_df['Predictions'].max() + 1)
    plt.grid()
    plt.savefig(os.path.join(shap_plot_dir, f'true_vs_predicted_number_{number}.png'))  # 保存图像
    plt.close()

# 打印预测结果
print("预测结果和真实值比较：")
print(comparison_df)
